#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <getopt.h>
#include <ctype.h>

#define OTA_HEADER_SIZE 32
#define MAX_FIRMWARE_TYPE_LENGTH 10

// Function prototypes
void print_help(const char *program_name);
int hex_to_bin(const char *input_file, const char *output_file);
int cly_to_bin(const char *input_file, const char *output_file);

int pack_ota_file(const char *input_file, const char *firmware_type, int major_version, int minor_version, const char *output_file) {
    // Determine the file format based on the extension
    const char *file_extension = strrchr(input_file, '.');
    if (file_extension == NULL) {
        fprintf(stderr, "Error: Invalid input file format.\n");
        return 0;
    }

    // Convert to binary if necessary
    char *temp_input_file = strdup(input_file);
    char temp_bin_file[256];
    snprintf(temp_bin_file, sizeof(temp_bin_file), "%s.bin", strtok(temp_input_file, "."));
    free(temp_input_file);

    if (strcmp(file_extension, ".hex") == 0) {
        if (!hex_to_bin(input_file, temp_bin_file)) {
            fprintf(stderr, "Error: Failed to convert .hex to .bin.\n");
            return 0;
        }
        input_file = temp_bin_file;
    } else if (strcmp(file_extension, ".cly") == 0) {
        if (!cly_to_bin(input_file, temp_bin_file)) {
            fprintf(stderr, "Error: Failed to convert .cly to .bin.\n");
            return 0;
        }
        input_file = temp_bin_file;
    }

    // Read the binary file
    FILE *bin_file = fopen(input_file, "rb");
    if (bin_file == NULL) {
        fprintf(stderr, "Error: Failed to open input file.\n");
        return 0;
    }

    // Get the size of the binary file
    fseek(bin_file, 0, SEEK_END);
    long bin_size = ftell(bin_file);
    rewind(bin_file);

    // Create the output file
    FILE *ota_file = fopen(output_file, "wb");
    if (ota_file == NULL) {
        fprintf(stderr, "Error: Failed to create output file.\n");
        fclose(bin_file);
        return 0;
    }

    // Write the OTA header
    unsigned char ota_header[OTA_HEADER_SIZE] = {0};
    ota_header[0] = 0xAA;  // 标识码1
    ota_header[1] = 0x55;  // 标识码2
    ota_header[2] = 1;     // 固件数量
    ota_header[3] = (unsigned char)major_version;  // 固件主版本号
    ota_header[4] = (unsigned char)minor_version;  // 固件次版本号
    ota_header[16] = (unsigned char)(bin_size & 0xFF);  // 固件长度 (低字节)
    ota_header[17] = (unsigned char)((bin_size >> 8) & 0xFF);
    ota_header[18] = (unsigned char)((bin_size >> 16) & 0xFF);
    ota_header[19] = (unsigned char)((bin_size >> 24) & 0xFF);

    // Calculate and write the checksum
    unsigned int checksum = 0;
    for (long i = 0; i < bin_size; i++) {
        unsigned char byte;
        fread(&byte, 1, 1, bin_file);
        checksum += byte;
    }
    ota_header[20] = (unsigned char)(checksum & 0xFF);  // 校验和 (低字节)
    ota_header[21] = (unsigned char)((checksum >> 8) & 0xFF);

    // Write the firmware type
    strncpy((char *)&ota_header[22], firmware_type, MAX_FIRMWARE_TYPE_LENGTH);

    fwrite(ota_header, 1, OTA_HEADER_SIZE, ota_file);

    // Write the binary data
    fseek(bin_file, 0, SEEK_SET);
    while (ftell(bin_file) < bin_size) {
        unsigned char buffer[1024];
        size_t bytes_read = fread(buffer, 1, sizeof(buffer), bin_file);
        fwrite(buffer, 1, bytes_read, ota_file);
    }

    // Close files
    fclose(bin_file);
    fclose(ota_file);

    // Remove temporary binary file if it exists
    if (strcmp(file_extension, ".hex") == 0 || strcmp(file_extension, ".cly") == 0) {
        remove(temp_bin_file);
    }

    return 1;
}
int unpack_ota_file(const char *input_file, const char *output_file) {
    // Open the input file
    FILE *ota_file = fopen(input_file, "rb");
    if (ota_file == NULL) {
        fprintf(stderr, "Error: Failed to open input file.\n");
        return 0;
    }

    // Read the OTA header
    unsigned char ota_header[OTA_HEADER_SIZE];
    if (fread(ota_header, 1, OTA_HEADER_SIZE, ota_file) != OTA_HEADER_SIZE) {
        fprintf(stderr, "Error: Invalid OTA file format.\n");
        fclose(ota_file);
        return 0;
    }

    // Verify the magic numbers
    if (ota_header[0] != 0xAA || ota_header[1] != 0x55) {
        fprintf(stderr, "Error: Invalid OTA file format (magic numbers mismatch).\n");
        fclose(ota_file);
        return 0;
    }

    // Extract the firmware length
    long firmware_length = (long)ota_header[16] + (long)(ota_header[17] << 8) + (long)(ota_header[18] << 16) + (long)(ota_header[19] << 24);

    // Create the output file
    FILE *bin_file = fopen(output_file, "wb");
    if (bin_file == NULL) {
        fprintf(stderr, "Error: Failed to create output file.\n");
        fclose(ota_file);
        return 0;
    }

    // Write the binary data
    unsigned char buffer[1024];
    while (firmware_length > 0) {
        size_t bytes_to_read = (firmware_length > sizeof(buffer)) ? sizeof(buffer) : firmware_length;
        size_t bytes_read = fread(buffer, 1, bytes_to_read, ota_file);
        fwrite(buffer, 1, bytes_read, bin_file);
        firmware_length -= bytes_read;
    }

    // Close files
    fclose(ota_file);
    fclose(bin_file);

    return 1;
}
int view_ota_info(const char *input_file) {
    // Open the input file
    FILE *ota_file = fopen(input_file, "rb");
    if (ota_file == NULL) {
        fprintf(stderr, "Error: Failed to open input file.\n");
        return 0;
    }

    // Read the OTA header
    unsigned char ota_header[OTA_HEADER_SIZE];
    if (fread(ota_header, 1, OTA_HEADER_SIZE, ota_file) != OTA_HEADER_SIZE) {
        fprintf(stderr, "Error: Invalid OTA file format.\n");
        fclose(ota_file);
        return 0;
    }

    // Verify the magic numbers
    if (ota_header[0] != 0xAA || ota_header[1] != 0x55) {
        fprintf(stderr, "Error: Invalid OTA file format (magic numbers mismatch).\n");
        fclose(ota_file);
        return 0;
    }

    // Extract and print the firmware type
    char firmware_type[MAX_FIRMWARE_TYPE_LENGTH + 1];
    strncpy(firmware_type, (char *)&ota_header[22], MAX_FIRMWARE_TYPE_LENGTH);
    firmware_type[MAX_FIRMWARE_TYPE_LENGTH] = '\0';
    printf("Firmware Type: %s\n", firmware_type);

    // Extract and print the firmware version
    int major_version = (int)ota_header[3];
    int minor_version = (int)ota_header[4];
    printf("Firmware Version: %d.%d\n", major_version, minor_version);

    // Extract and print the firmware length
    long firmware_length = (long)ota_header[16] + (long)(ota_header[17] << 8) + (long)(ota_header[18] << 16) + (long)(ota_header[19] << 24);
    printf("Firmware Length: %ld bytes\n", firmware_length);

    // Close the file
    fclose(ota_file);

    return 1;
}
int hex_to_bin(const char *input_file, const char *output_file) {
    FILE *hex_file = fopen(input_file, "r");
    if (hex_file == NULL) {
        fprintf(stderr, "Error: Failed to open input file.\n");
        return 0;
    }

    FILE *bin_file = fopen(output_file, "wb");
    if (bin_file == NULL) {
        fprintf(stderr, "Error: Failed to create output file.\n");
        fclose(hex_file);
        return 0;
    }

    char line[256];
    while (fgets(line, sizeof(line), hex_file)) {
        // Skip lines that do not start with ':'
        if (line[0] != ':') {
            continue;
        }

        // Parse the record length
        int record_length = (int)strtol(&line[1], NULL, 16);

        // Parse the address
        unsigned int address = (unsigned int)strtol(&line[3], NULL, 16);

        // Parse the record type
        int record_type = (int)strtol(&line[7], NULL, 16);

        // Calculate the checksum
        unsigned char checksum = (unsigned char)(~((record_length + (address >> 8) + (address & 0xFF) + record_type) + 0x01) + 1);

        // Verify the checksum
        for (int i = 9; i < 9 + 2 * record_length; i += 2) {
            unsigned char byte = (unsigned char)strtol(&line[i], NULL, 16);
            checksum += byte;
        }
        if (checksum != 0) {
            fprintf(stderr, "Error: Checksum mismatch in Intel Hex file.\n");
            fclose(hex_file);
            fclose(bin_file);
            remove(output_file);
            return 0;
        }

        // Write the data to the binary file
        for (int i = 9; i < 9 + 2 * record_length; i += 2) {
            unsigned char byte = (unsigned char)strtol(&line[i], NULL, 16);
            fwrite(&byte, 1, 1, bin_file);
        }
    }

    fclose(hex_file);
    fclose(bin_file);
    return 1;
}
int cly_to_bin(const char *input_file, const char *output_file) {
    FILE *cly_file = fopen(input_file, "r");
    if (cly_file == NULL) {
        fprintf(stderr, "Error: Failed to open input file.\n");
        return 0;
    }

    FILE *bin_file = fopen(output_file, "wb");
    if (bin_file == NULL) {
        fprintf(stderr, "Error: Failed to create output file.\n");
        fclose(cly_file);
        return 0;
    }

    // Read the first line to get the file length
    char line[256];
    if (fgets(line, sizeof(line), cly_file) == NULL) {
        fprintf(stderr, "Error: Invalid .cly file format.\n");
        fclose(cly_file);
        fclose(bin_file);
        remove(output_file);
        return 0;
    }

    // Parse the file length
    unsigned int file_length = (unsigned int)strtol(line, NULL, 16);

    // Read and convert the rest of the file
    while (fgets(line, sizeof(line), cly_file)) {
        for (int i = 0; i < 256 && line[i] != '\0'; i += 2) {
            unsigned char byte = (unsigned char)strtol(&line[i], NULL, 16);
            fwrite(&byte, 1, 1, bin_file);
        }
    }

    fclose(cly_file);
    fclose(bin_file);
    return 1;
}

int main(int argc, char *argv[]) {
    int opt;
    while ((opt = getopt(argc, argv, "hp:u:v:")) != -1) {
        switch (opt) {
            case 'h':
                print_help(argv[0]);
                return 0;
case 'p': {
                if (optind + 2 >= argc) {
                    fprintf(stderr, "Error: Missing required arguments for -p option.\n");
                    return 1;
                }
                const char *file = optarg;
                const char *type = argv[optind];
                const char *version = argv[optind + 1];
                // Validate and parse version
                int major_version, minor_version;
                if (sscanf(version, "%d.%d", &major_version, &minor_version) != 2 || major_version < 0 || major_version > 255 || minor_version < 0 || minor_version > 255) {
                    fprintf(stderr, "Error: Invalid version format. Expected a.b where 0 <= a,b <= 255.\n");
                    return 1;
                }
                // Determine output file name
                char *temp_file = strdup(file);
                char output_file[256];
                snprintf(output_file, sizeof(output_file), "%s_ota.bin", strtok(temp_file, "."));
                free(temp_file);
                if (pack_ota_file(file, type, major_version, minor_version, output_file)) {
                    printf("Packed %s to %s\n", file, output_file);
                } else {
                    return 1;
                }
                break;
            }
            case 'u':
                if (unpack_ota_file(optarg, "unpacked.bin")) {
                    printf("Unpacked %s to unpacked.bin\n", optarg);
                } else {
                    return 1;
                }
                break;
            case 'v':
                if (view_ota_info(optarg)) {
                    printf("Viewed info of %s\n", optarg);
                } else {
                    return 1;
                }
                break;
            default:
                fprintf(stderr, "Unknown option: -%c\n", opt);
                print_help(argv[0]);
                return 1;
        }
    }
    return 0;
}

void print_help(const char *program_name) {
    printf("Usage: %s [options]\n", program_name);
    printf("Options:\n");
    printf("  -h       Show help message\n");
    printf("  -p [file] [type] [version]  Pack firmware image to OTA update package\n");
    printf("  -u [file]              Unpack OTA update package to binary image\n");
    printf("  -v [file]              View information of OTA update file\n");
    printf("\nSupported file formats: .hex, .bin, .cly\n");
    printf("Supported firmware types: Any ASCII string up to 10 bytes\n");
    printf("Version format: a.b where 0 <= a,b <= 255 (a is major version, b is minor version)\n");
}

// Implementations of the above functions will be added in subsequent steps.
